#!/usr/bin/env python
# encoding: utf-8
"""
core.py

Foundation classes for the ODEPY library.

ODEPY is a Python module which consists of a class library 
for time stepping integrators for ordinary differential
equations. It is meant to facilitate the construction and 
analysis of numerical integrators for initial value problems.

The module requires a full installation of Numpy/Scipy.

The code is available under GNU GPL v3 licence at
http://code.google.com/p/odepy

Created by Klas Modin on 2010-06-01.
"""

__all__ = ["Solver","OneStepMethod","Composition"]

from numpy import array, dot, ones, zeros, hstack
from numpy import diag, triu, isscalar, linspace, shape

class Solver(object):	
	"""
	Abstract class for autonomous ODE solver.
	
	Attributes
	----------
	name : string
		Name of the solver.
	
	Methods
	-------
	integrate
		Carry out the integration process.
	"""
	
	name = 'Solver'
	
	def __init__(self,name=None):
		"""
		Constructor. 
		
		Parameters
		----------
		name : string, optional
			Name of the solver.
		"""
		if name:
			self.name = name
	
	def integrate(self, tspan, yinit):
		"""
		Carry out the integration process. Subclasses must
		implement this method.
		
		Parameters
		----------
		tspan : float
			Total integration time
		
		yinit : array_like
			Initial conditions.
		
		Returns
		-------
		tout : ndarray, shape (M,)
			Vector with output times.
		yout : ndarray
			Array with states corresponding to output times.
		"""
		raise NotImplementedError('integrate method for class %s is not defined.'%str(type(self)))
	
	def __str__(self):
		try:
			try:
				funstr = '('+self.f.func_name+')'
			except AttributeError:
				funstr = ''
			return self.name+funstr
		except AttributeError:
			return super(Solver, self).__str__()

class OneStepMethod(Solver):
	
	name = 'OneStep'
	
	def stepwise(self, tspan, yinit, stepsize):
		"""
		Carry out the integration process step-wise.
		"""
		if isscalar(yinit):
			y = asarray(yinit)
		else:
			y = array(yinit)
		nout = int(round(float(tspan)/stepsize))
		yield y
		for i in xrange(1,nout):
			y = self.__call__(stepsize,y)
			yield y
	
	def integrate(self, tspan, yinit, stepsize):
		"""
		Carry out the integration process. Subclasses must
		implement this method.
		
		Parameters
		----------
		tspan : float
			Total integration time.
		
		yinit : array_like or scalar
			Initial conditions.
		
		stepsize : float
			The time increment step size.
		
		Returns
		-------
		tout : ndarray, shape (M,)
			Vector with output times.
		yout : ndarray
			Array with states corresponding to output times.
		"""
		if isscalar(yinit):
			y = asarray(yinit)
		else:
			y = array(yinit)
		nout = int(round(float(tspan)/stepsize))
		tout = linspace(0,tspan,nout)
		yout = zeros(hstack((nout,shape(y))),dtype=y.dtype)
		yout[0] = y
		for i in xrange(1,nout):
			# self.__call__(stepsize,y)
			y = self.__call__(stepsize,y)
			yout[i] = y
			# yout[i] = self.__call__(stepsize,yout[i-1])
		
		return (tout,yout)
	
	def __call__(self,h,y):
		"""
		This is evaluation of the numerical flow function.
		Take a step with step size 'h' from input phase data 'y'
		and return new output phase vector.
		Notice that 'y' is used as internal storage,
		so it is overwritten.
		"""
		raise NotImplementedError('__call__ method for class %s is not defined.'%str(type(self)))
	
	def __mul__(self,other):
		"""
		Return composition of two methods.
		
		Parameters
		----------
		other : OneStepMethod
			The method that self will be composed with.
			Returning method is: (h,y) -> self(h,.) o other(h,y).
		"""
		if (not isinstance(other,OneStepMethod)):
			raise AttributeError("Object %s must be a OneStepMethod."%str(other))
		if (isinstance(self,Composition)):
			mclist_self = self.mclist
		else:
			mclist_self = [(self,1.0)]
		if (isinstance(other,Composition)):
			mclist_other = other.mclist
		else:
			mclist_other = [(other,1.0)]
		return Composition(mclist_self+mclist_other)
	
	def __pow__(self,coeff):
		"""
		Return method with "stepsize coefficient".
		
		Parameters
		----------
		coeff : scalar
			The coefficient. Returned method is (h,y) -> self(coeff*h,y).
		"""
		if (not isscalar(coeff)):
			raise AttributeError("Object %s must be a scalar."%str(coeff))
		return Composition([(self,coeff)])
	

class Composition(OneStepMethod):
	"""
	Composition of two or more methods.
	"""
	name = "Composition"
	
	def __init__(self, mclist):
		"""
		Constructor.
		
		Parameters
		----------
		mlist : list of tuples (OneStepMethod,coeff)
			List of tuples. Each tuple should contains a
			method, and the coefficient for that method.
		"""
		super(Composition, self).__init__()
		self.mclist = mclist
	
	def __call__(self,h,y):
		for (m,c) in self.mclist:
			y = m.__call__(c*h,y)
		return y
	
	def __pow__(self,coeff):
		if (not isscalar(coeff)):
			raise AttributeError("Object %s must be a scalar."%str(coeff))
		mclist_new = []
		for (m,c) in self.mclist:
			mclist_new += [(m,coeff*c)]
		return Composition(mclist_new)
	
	def __str__(self):
		namestr = ''
		for (m,c) in self.mclist:
			if c == 1.0:
				namestr += str(m)+' * '
			else:
				namestr += "%s**%s * "%(str(m),str(c))
		return namestr[:-3]
	


class ImplicitOneStep(OneStepMethod):
	"""
	Class for implicit one-step methods.
	
	Attributes
	----------
	A : ndarray, shape (s,s)
		The A matrix for the method, where s is the number of stages.
	
	b : ndarray, shape (s,)
		The b vector for the method.
	
	f : function
		The RHS of the equation. Should have the form
			"def f(y)"
		and should return an ndarray of same shape as y.
		The input y should not be overwritten by f.
	
	jac : function
		The Jacobian corresponding to the RHS. Should
		have the form
			"def jac(y)"
		and should return an ndarray one order higher
		than that of y. The input y should not be
		overwritten by jac.
	
	
	
	params : dict
		Parameters for tolerances etc.
	
	stats : dict
		Statistics about total number of steps, Jacobians,
		rhs evaluations etc.
	
	Methods
	-------
	integrate
		Carry out the integration process.
	"""
	
	name = "Implicit one-step"
	
	def __init__(self, f, A, b, jac=None, params=dict()):
		"""
		Constructor.
		
		Parameters
		----------
		f : function
			The RHS of the equation. Should have the form
				"def f(y)"
			and should return an ndarray of same shape as y.
		
		A : ndarray, shape (s,s)
			The A matrix in the Butcher tableau for the method, 
			where s is the number of stages.
			Matrix should be strictly lower triangular 
			(so that the method is explicit).
		
		b : ndarray, shape (s,)
			The b vector in the Butcher tableau for the method.
		
		jac : function (optional)
			The Jacobian corresponding to the RHS. Should
			have the form
				"def jac(y)"
			and should return an ndarray one order higher
			than that of y. The input y should not be
			overwritten by jac. If not jac is given, a
			finite difference approximation will be used.
		
		params : dict (optional)
			Parameters for the root solving:
			- newtol (real), Newton tolerance
			- kmax (int), maximal number of iterations
			- kappa (real), safety factor
		"""
		super(ImplicitRungeKutta, self).__init__()
		dims = A.shape
		if (dims[0] != 1 and dims[0] != dims[1]):
			raise AttributeError("A must be a square matrix.")
		s = dims[0]
		self.A = array(A)
		if (len(b) != s):
			raise AttributeError("Size of matrix A and vector b does not match.")
		self.b = array(b)
		self.f = f
		if jac is None:
			# Notice: this Jacobian only works in the case when y is
			# a one-dimensional array.
			def jac(y):
				f0 = f(y)
				# ynew = y.copy()
				delta = 1e-10
				# for i,v in ndenumerate(ynew):
				# 	ynew(i) += delta
				J = identity(len(y))
				for Jk in J:
					Jk[:] = (f(y+delta*Jk)-f0)/delta
				return J.T
			
		self.jac = jac
		
		# Setup parameters
		self.params = dict()
		try:
			self.params['newtol'] = params['newtol']
		except KeyError:
			self.params['newtol'] = sqrt(finfo(1.0).eps) # Sqrt of machine epsilon (for default float type).
		try:
			self.params['kmax'] = params['kmax']
		except KeyError:
			self.params['kmax'] = 7
		try:
			self.params['kappa'] = params['kappa']
		except KeyError:
			self.params['kappa'] = 1e-1
		
		# Setup stats
		self.stats = dict()
		self.stats['njac'] = 0
		self.stats['nlu'] = 0
		self.stats['nrhs'] = 0
		self.stats['nsteps'] = 0
		self.stats['niterations'] = 0
		
		
	
	def initial_setup(self,y):
		# Dimension variables
		s = len(self.b)
		n = len(y)
		
		self.hlast = 0
		self.updateJacobian = True
		self.indices = arange(s*n).reshape((s,n))
		self.Z = zeros(s*n)
		self.dZ = zeros_like(self.Z)
		self.eta = 0
		self.Uround = finfo(0.0).eps
		if (det(self.A) != 0.0):
			self.d = dot(self.b,inv(self.A))
	
	def __call__(self,h,y):
		
		# Check for first call, if so, run setup.
		if (hasattr(self,'hlast') == False):
			self.initial_setup(y)
		
		# 1. Update data before iteration.
		self.iterationStepsTaken = 0
		self.update_before_iteration()
		
		# 2. Predict the variable Z (prediction step).
		self.Z[:] = self.predict_Z()
		
		# 3. First evaluation of F at 
		
		# Dimension variables
		s = len(self.b)
		n = len(y)
		
		# Check what needs to be updated:
		if (self.hlast != h):
			self.updateLU = True
			self.hlast = h
		
		if self.updateJacobian:
			self.J = self.jac(y)
			self.updateLU = True
			self.stats['njac'] += 1
		
		if self.updateLU:
			(self.lu,self.piv) = lu_factor(identity(s*n)-h*kron(self.A,self.J))
			self.stats['nlu'] += 1
			self.updateLU = False
		
		# Set initial Z
		self.Z[:] = 0.0
		
		# First evaluation of F at Z
		self.F = kron(ones(s),self.f(y))
		self.stats['nrhs'] += 1
		
		# Iterate
		for k in xrange(self.params['kmax']):
			
			# Update statistics
			self.stats['niterations'] += 1
			
			# One Newton step
			self.dZ = lu_solve((self.lu,self.piv),h*dot(kron(self.A,identity(n)),self.F)-self.Z)
			
			# Update Z
			self.Z += self.dZ
			
			# Check stopping criterion
			if (k==0):
				self.eta = max(self.eta,self.Uround)**0.8
				self.dZnorm = norm(self.dZ)
			else:
				dZnormlast = self.dZnorm
				self.dZnorm = norm(self.dZ)
				Theta = self.dZnorm/dZnormlast
				if Theta > 1:
					raise AssertionError("The iteration is diverging with Theta=%f"%Theta)
				self.eta = Theta/(1-Theta)
			if (self.eta*self.dZnorm <= self.params['kappa']*self.params['newtol']):
				break
			
			# Evaluate F
			for i in self.indices:
				self.F[i] = self.f(y+self.Z[i])
				self.stats['nrhs'] += 1
			
		else:
			raise AssertionError("Iteration did not converge in %i iterations."%(k+1))
		
		# Update phase space variable
		if (hasattr(self,'d')):
			y += dot(self.d,self.Z[self.indices])
		else:
			# (I-h*AoJ)(Z1-Z0) = -Z0 + h AoI F(Z0)
			# Z1 = h*AoI F(Z1)
			# y1 = y0 + h*dot(b,F(Z1))
			
			# Evaluate F
			for i in self.indices:
				self.F[i] = self.f(y+self.Z[i])
				self.stats['nrhs'] += 1
			y += h*dot(self.b,self.F[self.indices])
			
		
		# Set what to do next step
		if (k==0 or Theta < 1e-3):
			self.updateJacobian = False
		else:
			self.updateJacobian = True
		
		# Update statistics
		self.stats['nsteps'] += 1
		self.stats['iterations/step'] = float(self.stats['niterations'])/self.stats['nsteps']
		self.stats['jac/step'] = float(self.stats['njac'])/self.stats['nsteps']
		self.stats['rhs/step'] = float(self.stats['nrhs'])/self.stats['nsteps']
		
		# Return output
		return y
	


class ImplicitMethod(OneStepMethod):
	"""
	Abstract class for implicit methods.
	Contains a Newton solver.
	"""
	name = "ImplicitMethod"
	
	def __init__(self, fun, funjac):
		super(ImplicitMethod, self).__init__()
		self.fun = fun
		self.funjac = funjac
		
	
	def __call__(self,h,y):
		ynew = y.copy()
		self.find_root(y0=y,y1=ynew,h=h)
		return ynew
	
	def find_root(self,y0,y1,h):
		"""Find root using Newton's method."""
		pass


if __name__ == '__main__' and __file__.split('/')[-1] == 'core.py':
	from numpy import sin
	# Setup and solve non-linear pendulum problem
	
	# Efficient implementation of the Störmer-Verlet scheme.
	class FlowA(OneStepMethod):
		name = 'FlowA'
		def __call__(self,h,y):
			dim = len(y)/2
			y[dim:] -= h*sin(y[:dim])
			return y
		
	
	class FlowB(OneStepMethod):
		name = 'FlowB'
		def __call__(self,h,y):
			dim = len(y)/2
			y[:dim] += h*y[dim:]
			return y
		
	
	# Störmer-Verlet
	mA = FlowA()
	mB = FlowB()
	solver = mA**0.5*mB*mA**0.5
	
	y0 = hstack((0,0.5))
	tend = 100.0
	stepsize = 0.5
	
	(tout,yout) = solver.integrate(tend,y0,stepsize)
	